#include "os.h"
#include <stddef.h>

extern ptr_t TEXT_START;
extern ptr_t TEXT_END;
extern ptr_t DATA_START;
extern ptr_t DATA_END;
extern ptr_t RODATA_START;
extern ptr_t RODATA_END;
extern ptr_t BSS_START;
extern ptr_t BSS_END;
extern ptr_t HEAP_START;
extern ptr_t HEAP_SIZE;

/*
 * _alloc_start points to the actual start address of heap pool
 * _alloc_end points to the actual end address of heap pool
 * _num_pages holds the actual max number of pages we can allocate.
 */
static ptr_t _alloc_start = 0;
static ptr_t _alloc_end = 0;
static uint32_t _num_pages = 0;

#define PAGE_SIZE 4096
#define PAGE_ORDER 12


/*
 * Page Descriptor 
 * flags:
 * - bit 0: flag if this page is taken(allocated)
 * - bit 1: flag if this page is the last page of the memory block allocated
 */
#define PAGE_TAKEN (uint8_t)(1 << 0)
#define PAGE_LAST  (uint8_t)(1 << 1)
struct Page {
	uint8_t flags;
};
/* inline内联函数，建议编译器在编译时将函数的定义直接插入
 * 到每个函数调用的地方，而不是进行常规的函数调用。
 * 这样可以减少函数调用的开销，提高程序的执行效率。
 * 但是，内联函数的使用也有一些限制和注意事项：
 * 1. 内联函数的代码会被插入到每个调用的地方，这可能会导致代码膨胀，增加可执行文件的大小。
 * 2. 内联函数不能包含循环、递归或复杂的控制结构，因为这些结构可能导致编译器无法确定内联的边界。
 * 3. 内联函数不能有静态变量，因为静态变量在每个调用中都需要保持状态，而内联函数的代码在每个调用中都是独立的。
 * 4. 内联函数不能有可变参数列表，因为编译器无法确定参数的数量和类型。
 */

static inline void _clear(struct Page *page)
{
	page->flags = 0;
}

static inline void _set_flag(struct Page *page, uint8_t flags)
{
	page->flags |= flags;
}

static inline int _is_free(struct Page *page)
{
	if (page->flags & PAGE_TAKEN) {
		return 0;
	} else {
		return 1;
	}
}

static inline int _is_last(struct Page *page)
{
	if (page->flags & PAGE_LAST) {
		return 1;
	} else {
		return 0;
	}
}

/*
 * align the address to the border of page(4K)
 */
static inline ptr_t _align_page(ptr_t address)
{
	ptr_t order = (1 << PAGE_ORDER) - 1;
	return (address + order) & (~order);
}

/*
 *    ______________________________HEAP_SIZE_______________________________
 *   /   ___num_reserved_pages___   ______________num_pages______________   \
 *  /   /                        \ /                                     \   \
 *  |---|<--Page-->|<--Page-->|...|<--Page-->|<--Page-->|......|<--Page-->|---|
 *  A   A                         A                                       A   A
 *  |   |                         |                                       |   |
 *  |   |                         |                                       |   _memory_end
 *  |   |                         |                                       |
 *  |   _heap_start_aligned       _alloc_start                            _alloc_end
 *  HEAP_START(BSS_END)
 *
 *  Note: _alloc_end may equal to _memory_end.
 */
void page_init()
{
	ptr_t _heap_start_aligned = _align_page(HEAP_START);

	/* 
	 * We reserved some Pages to hold the Page structures.
	 * The number of reserved pages depends on the LENGTH_RAM.
	 * For simplicity, the space we reserve here is just an approximation,
	 * assuming that it can accommodate the maximum LENGTH_RAM.
	 * We assume LENGTH_RAM should not be too small, ideally no less
	 * than 16M (i.e. PAGE_SIZE * PAGE_SIZE).
	 */
	uint32_t num_reserved_pages = LENGTH_RAM / (PAGE_SIZE * PAGE_SIZE);

	_num_pages = (HEAP_SIZE - (_heap_start_aligned - HEAP_START))/ PAGE_SIZE - num_reserved_pages;
	printf("HEAP_START = %p(aligned to %p), HEAP_SIZE = 0x%lx,\n"
	       "num of reserved pages = %d, num of pages to be allocated for heap = %d\n",
	       HEAP_START, _heap_start_aligned, HEAP_SIZE,
	       num_reserved_pages, _num_pages);
	
	/*
	 * We use HEAP_START, not _heap_start_aligned as begin address for
	 * allocating struct Page, because we have no requirement of alignment
	 * for position of struct Page.
	 * The index starts at HEAP_START, and each struct Page represents the actual memory page
	 * 初始化范围为HEAP_START---> num_pages * sizeof(struct Page)
	 */
	struct Page *page = (struct Page *)HEAP_START;
	for (int i = 0; i < _num_pages; i++) {
		_clear(page);
		page++;	
	}

	_alloc_start = _heap_start_aligned + num_reserved_pages * PAGE_SIZE;
	_alloc_end = _alloc_start + (PAGE_SIZE * _num_pages);

	printf("TEXT:   %p -> %p\n", TEXT_START, TEXT_END);
	printf("RODATA: %p -> %p\n", RODATA_START, RODATA_END);
	printf("DATA:   %p -> %p\n", DATA_START, DATA_END);
	printf("BSS:    %p -> %p\n", BSS_START, BSS_END);
	printf("HEAP:   %p -> %p\n", _alloc_start, _alloc_end);
}

/*
 * Allocate a memory block which is composed of contiguous physical pages
 * - npages: the number of PAGE_SIZE pages to allocate
 */
void *page_alloc(int npages)
{
	/* Note we are searching the page descriptor bitmaps. */
	int found = 0;
	int j = 0;
	struct Page *page_i = (struct Page *)HEAP_START;
	for (int i = 0; i <= (_num_pages - npages); i++) {
		if (_is_free(page_i)) {
			found = 1;
			/* 
			 * meet a free page, continue to check if following
			 * (npages - 1) pages are also unallocated.
			 */
			struct Page *page_j = page_i + 1;
			for (j = i + 1; j < (i + npages); j++) {
				if (!_is_free(page_j)) {
					found = 0;
					break;
				}
				page_j++;
			}
			/*
			 * get a memory block which is good enough for us,
			 * take housekeeping, then return the actual start
			 * address of the first page of this memory block
			 */
			if (found) {
				struct Page *page_k = page_i;
				for (int k = i; k < (i + npages); k++) {
					_set_flag(page_k, PAGE_TAKEN);
					page_k++;
				}
				page_k--;
				_set_flag(page_k, PAGE_LAST);
				return (void *)(_alloc_start + i * PAGE_SIZE);
			}
		}
		page_i += (j + 1);
	}
	return NULL;
}

/*
 * Free the memory block
 * - p: start address of the memory block
 */
void page_free(void *p)
{
	/*
	 * Assert (TBD) if p is invalid
	 */
	if (!p || (ptr_t)p >= _alloc_end) {
		return;
	}
	/* get the first page descriptor of this memory block */
	struct Page *page = (struct Page *)HEAP_START;
	page += ((ptr_t)p - _alloc_start)/ PAGE_SIZE;
	/* loop and clear all the page descriptors of the memory block */
	while (!_is_free(page)) {
		if (_is_last(page)) {
			_clear(page);
			break;
		} else {
			_clear(page);
			page++;;
		}
	}
}

void page_test(void)
{
	void *p = page_alloc(2);
	printf("p = %p\n", p);
	//page_free(p);

	void *p2 = page_alloc(7);
	printf("p2 = %p\n", p2);
	page_free(p2);

	void *p3 = page_alloc(4);
	printf("p3 = %p\n", p3);
}


/* 
 * 字节为单位的内存管理策略：
 * 该策略通过将内存划分为固定大小的块，
 * 并使用链表来管理这些块的分配和释放。
 * 每个块包含一个头部（BlockHeader），用于记录块的大小、是否空闲以及指向下一个块的指针。
 * 内存分配时，会从空闲块链表中查找合适的块，并在必要时分割块以满足请求。
 * 释放内存时，会将块标记为“空闲”，并尝试合并相邻的空闲块以减少内存碎片。
 * 整个内存管理以字节为单位进行操作，确保内存的高效利用。
 */

#define ALIGNMENT 8 			
#define ALIGN(size) (((size) + (ALIGNMENT - 1)) & ~(ALIGNMENT - 1))

typedef struct BlockHeader {
    size_t size;          		// 块大小（包括头部）
    struct BlockHeader *next; 	// 指向下一个块
    uint8_t free;              	// 是否空闲
} BlockHeader;

static BlockHeader *heap_start = NULL; // 堆的起始地址
static BlockHeader *free_list = NULL;  // 空闲块链表

// 初始化堆
static void heap_init() {
    if (!heap_start) {
		// 分配一个初始页面
        heap_start = (BlockHeader *)page_alloc(1); 
        if (!heap_start) {
            printf("Heap initialization failed!\n");
            return;
        }
        heap_start->size = PAGE_SIZE;
        heap_start->next = NULL;
        heap_start->free = 1;
        free_list = heap_start;
    }
}

// 分配内存
void *malloc(size_t size) {
    if (size == 0) {
        return NULL;
    }
	/* 整个内存块（头+数据）按8字节对齐 */
    size = ALIGN(size + sizeof(BlockHeader)); 

    if (!heap_start) {
        heap_init();
    }

    BlockHeader *current = free_list;
    BlockHeader *prev = NULL;

    // 首次适配算法
    while (current) {
        if (current->free && current->size >= size) {
            // 如果块足够大，分割块
            if (current->size > size + sizeof(BlockHeader)) {
                BlockHeader *new_block = (BlockHeader *)((char *)current + size);
                new_block->size = current->size - size;
                new_block->free = 1;
                new_block->next = current->next;

                current->size = size;
                current->next = new_block;
            }
			// 如果当前块的大小刚好等于所需大小，则无需分割
            current->free = 0;
            return (void *)((char *)current + sizeof(BlockHeader));
        }
        prev = current;
        current = current->next;
    }

    // 如果没有找到合适的块，分配新的页面
    BlockHeader *new_block = (BlockHeader *)page_alloc(1);
    if (!new_block) {
        return NULL;
    }
    new_block->size = PAGE_SIZE;
    new_block->free = 0;
    new_block->next = NULL;

    if (prev) {
        prev->next = new_block;
    }

    return (void *)((char *)new_block + sizeof(BlockHeader));
}

// 释放内存
void free(void *ptr) {
    if (!ptr) {
        return;
    }

    BlockHeader *block = (BlockHeader *)((char *)ptr - sizeof(BlockHeader));
    block->free = 1;

    // 合并相邻的空闲块
    BlockHeader *current = free_list;
    while (current) {
        if (current->free && current->next && current->next->free) {
            current->size += current->next->size;
            current->next = current->next->next;
        }
        current = current->next;
    }
}

// 测试 malloc 和 free
void malloc_test(void) {
    void *p1 = malloc(100);
    printf("Allocated p1 = %p\n", p1);

    void *p2 = malloc(200);
    printf("Allocated p2 = %p\n", p2);

    free(p1);
    printf("Freed p1\n");

    void *p3 = malloc(50);
    printf("Allocated p3 = %p\n", p3);

    free(p2);
    printf("Freed p2\n");

    free(p3);
    printf("Freed p3\n");
}

